# -*- coding:utf8 -*-

import os
import sys
import random
import PIL.Image as Image

def walk_through_folder_for_split(src_folder):
    test_set  = []
    train_set = []

    label = 0
    flag = 1
    for tea_folder in os.listdir(src_folder):
        tea_path = src_folder +'\\'+ tea_folder
        tea_imgs = []
        for img_file in os.listdir(tea_path):
            img_path = tea_path +'\\'+ img_file
            im = Image.open(img_path)
            x_size, y_size = im.size
            if x_size==64 and y_size==64:
                tea_imgs.append((img_path, label))
                flag += 1
        random.shuffle(tea_imgs)
        if tea_folder=='negtive_all_64_aug':
            train_set += tea_imgs[0:26000]
            test_set  += tea_imgs[26000:len(tea_imgs)]
        elif tea_folder=='positive_all_64_aug':
            train_set += tea_imgs[0:5500]
            test_set  +=tea_imgs[5500:len(tea_imgs)]
        sys.stdout.flush()
        label += 1
    print 'all images:' + ''+ str(flag)
    print 'test  set num: %d' % (len(test_set))
    print 'train set num: %d' % (len(train_set))
    return test_set, train_set

def set_to_csv_file(data_set, file_name):
    f = open(file_name, 'wb')
    for item in data_set:
        line = item[0] + ',' + str(item[1]) + '\n'
        f.write(line)
    f.close()

if __name__ == '__main__':
    """
    test_set_file  = "C:\\Users\\wuxiaomin\\Desktop\\arobei\\testfile_160_aug"
    train_set_file = "C:\\Users\\wuxiaomin\\Desktop\\arobei\\trainfile_160_aug"
    test_set, train_set = walk_through_folder_for_split("C:\\Users\\wuxiaomin\\Desktop\\arobei\\tea_crop_160_aug")
    set_to_csv_file(test_set,  test_set_file)
    set_to_csv_file(train_set, train_set_file)
    """
    test_set_file  = "C:\\Users\\wuxiaomin\\Desktop\\arobei\\new\\testfile"
    train_set_file = "C:\\Users\\wuxiaomin\\Desktop\\arobei\\new\\trainfile"
    test_set, train_set = walk_through_folder_for_split("C:\\Users\\wuxiaomin\\Desktop\\arobei\\new\\data")
    set_to_csv_file(test_set,  test_set_file)
    set_to_csv_file(train_set, train_set_file)